"""
https://leetcode.cn/problems/classes-with-at-least-5-students/description/?envType=study-plan-v2&envId=30-days-of-pandas&lang=pythondata

596. 超过 5 名学生的课
已解答
简单
相关标签
premium lock icon
相关企业
SQL Schema
Pandas Schema
表: Courses

+-------------+---------+
| Column Name | Type    |
+-------------+---------+
| student     | varchar |
| class       | varchar |
+-------------+---------+
(student, class)是该表的主键（不同值的列的组合）。
该表的每一行表示学生的名字和他们注册的班级。
 

查询 至少有 5 个学生 的所有班级。

以 任意顺序 返回结果表。

结果格式如下所示。

 

示例 1:

输入: 
Courses table:
+---------+----------+
| student | class    |
+---------+----------+
| A       | Math     |
| B       | English  |
| C       | Math     |
| D       | Biology  |
| E       | Math     |
| F       | Computer |
| G       | Math     |
| H       | Math     |
| I       | Math     |
+---------+----------+
输出: 
+---------+ 
| class   | 
+---------+ 
| Math    | 
+---------+
解释: 
-数学课有 6 个学生，所以我们包括它。
-英语课有 1 名学生，所以我们不包括它。
-生物课有 1 名学生，所以我们不包括它。
-计算机课有 1 个学生，所以我们不包括它。

"""

import pandas as pd

def find_classes(courses: pd.DataFrame) -> pd.DataFrame:
    temp= courses.groupby('class').agg({
        'student':'size'
    }).reset_index()
    return temp.loc[temp['student']>=5].drop('student',axis=1)
    pass


if __name__=='__main__':
    data = {
        'student': ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I'],
        'class': ['Math', 'English', 'Math', 'Biology', 'Math', 'Computer', 'Math', 'Math', 'Math']
    }
    courses = pd.DataFrame(data)
    result = find_classes(courses)
    print(result)

